from datetime import datetime
import backtrader as bt
import akshare as ak
import pandas as pd
import numpy as np
import quantstats as qs
import webbrowser
from sqlalchemy import create_engine
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser
from collections import defaultdict
import warnings

warnings.filterwarnings('ignore')

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


def get_trade_info():
    sql1 = f"SELECT * FROM TKH_KDL_BT"
    sql2 = "SELECT STK_CODE, RIGHT(F_INDEX_CODE, 6) FROM ETF_INDEX_INFO"
    r1 = exec_sql(sql1)
    r2 = exec_sql(sql2)
    df1 = pd.DataFrame(r1, columns=['代码', 'datetime'])
    df2 = pd.DataFrame(r2, columns=['code', '代码'])
    df = pd.merge(df1, df2, on=['代码'])[['datetime', 'code', '代码']]
    etf_df = ak.fund_etf_spot_em().loc[:, ['代码', '名称', '流通市值']]
    code_list = list()

    for index in df['代码'].unique():
        codes = df.loc[df['代码'] == index, 'code'].to_list()
        ret = etf_df.loc[etf_df['代码'].isin(codes)]
        code = ret.loc[ret['流通市值'].idxmax(), '代码']
        code_list.append(code)

    df = df.loc[df['code'].isin(code_list), ['datetime', 'code']]
    dic = defaultdict(list)
    stocks = set()

    for i, row in df.iterrows():
        dt, code = row['datetime'], row['code']
        dic[dt].append(code)
        stocks.add(code)

    res = {'trade_code': dic, 'stocks': list(stocks)}
    return res


def get_data(stocks, freq='day', source='db'):
    if freq == '1min':
        table_name = 'STOCK_MIN_DATA'
    elif freq == '15min':
        table_name = 'STOCK_15MIN_DATA'
    else:
        table_name = 'ETF_DAY_DATA'  # HIS_DAILY_DATA

    res = dict()
    date_data = set_date()
    st = date_data['st']
    et = date_data['et']
    st_d = date_data['st_d']
    et_d = date_data['et_d']

    if source == 'db':
        print('db')
        stocks_str = ','.join(f"'{tok}'" for tok in stocks)
        sql = f"SELECT code, datetime, [close], [open], low, high, volume FROM {table_name} WHERE code in ({stocks_str})"
        r = exec_sql(sql)
        df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume'])
        # df['datetime'] = df['datetime'].map(time_map)
        # print(df)

        for code in stocks:
            if code == '399300':
                sig_df = ak.stock_zh_index_daily(symbol="sz399300")
                sig_df = sig_df.loc[(sig_df.date >= st_d) & (sig_df.date <= et_d), :]
                sig_df.rename(columns={'date': 'datetime'}, inplace=True)
                sig_df.insert(0, 'code', code)
                sig_df = sig_df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
            else:
                sig_df = df.loc[df['code'] == code]
            if len(sig_df) == 0:
                continue
            sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            sig_df.sort_values('datetime', inplace=True, ignore_index=True)
            sig_df.set_index('datetime', drop=False, inplace=True)
            print(f'code={code}->len={len(sig_df)}')
            # print(sig_df)
            res[code] = sig_df
    else:
        print('internet')
        for code in stocks:
            if code == '399300':
                df = ak.stock_zh_index_daily(symbol="sz399300")
                df = df.loc[(df.date >= st_d) & (df.date <= et_d), :]
                df.rename(columns={'date': 'datetime'}, inplace=True)
            else:
                df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=st, end_date=et, adjust="qfq")
                df.rename(columns={'日期': 'datetime', '开盘': 'open', '收盘': 'close', '最高': 'high',
                                   '最低': 'low', '成交量': 'volume'}, inplace=True)

            df.insert(0, 'code', code)
            df = df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
            df['datetime'] = pd.to_datetime(df['datetime'])
            df.set_index('datetime', inplace=True)
            print(f'code={code}->len={len(df)}')
            res[code] = df

    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def time_map(x):
    t = int(datetime.strftime(datetime.strptime(x, '%Y-%m-%d'), '%Y%m%d'))
    # t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


def set_date():
    st = "20190508"
    et = "20230816"
    st_dt = datetime.strptime(st, '%Y%m%d')
    et_dt = datetime.strptime(et, '%Y%m%d')
    st_d = st_dt.date()
    et_d = et_dt.date()
    res = {'st': st, 'et': et, 'st_d': st_d, 'et_d': et_d, 'st_dt': st_dt, 'et_dt': et_dt}

    return res


class TKH_KDL(bt.Strategy):
    params = (
        ('money', 1e4),
        ('maperiod', 26),
        ('times', 1),
        ('printlog', False),
    )

    def __init__(self):
        self.order = None
        self.trade_info = list()
        self.trade_code = trade_code
        self.trail_dict = dict()
        self.TR = dict()
        self.ATR = dict()
        self.daily_pnl = dict()

        for data in self.datas:
            code = data._name
            if code != '399300':
                self.trail_dict[code] = {'cdt': 0, 'cost': 0}
                self.TR[code] = bt.indicators.Max((data.high - data.low), abs(data.close(-1) - data.high),
                                                        abs(data.close(-1) - data.low))
                self.ATR[code] = bt.indicators.SMA(self.TR[code], period=self.p.maperiod, subplot=False)

    def log(self, txt, dt=None, doprint=True):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            content = f'{dt.isoformat()}, {txt}'
            print(content)

    def prenext(self):
        self.next()

    def downcast(self, price):
        res = self.p.money / price // 100 * 100
        return res

    def next(self):
        if self.order:
            return

        cur_dt = self.datas[0].datetime.datetime(0)
        dt_int = int(cur_dt.strftime('%Y%m%d'))

        if dt_int not in self.daily_pnl:
            self.daily_pnl[dt_int] = round(self.broker.getvalue() - init_cash, 2)

        for data in self.datas:
            code = data._name
            try:
                dt = data.datetime.datetime(0)  # 获取当前的回测时间点
            except:
                print(f'{code}-{dt}')
            if code == '399300' or dt != cur_dt:
                continue

            pos = self.getposition(data).size

            if dt_int in trade_code and code in trade_code[dt_int]:
                amount = self.downcast(data.close[0])
                self.log(f'code={code}->amount={amount}')
                if amount > 0:
                    self.order = self.buy(data=data, size=amount, name=code)
            elif pos and data.high[0] / self.trail_dict[code]['cost'] - 1 > 0.42:
                amount = pos
                self.order = self.sell(data=data, size=amount, name=code)

    def notify_order(self, order):
        code = order.data._name

        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:

            if order.isbuy():
                self.log(
                    f'code={code} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={code} BUY EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f'Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')

                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'BUY', round(order.executed.price, 2), order.executed.size,
                              round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
                self.trail_dict[code]['cdt'] = int(bt.num2date(order.executed.dt).strftime('%Y%m%d'))
                self.trail_dict[code]['cost'] = round(order.executed.price, 2)
            else:
                self.log(
                    f'code={order.data._name} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={order.data._name} SELL EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f' Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'SELL', round(order.executed.price, 2),
                              order.executed.size, round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f'Order Canceled/Margin/Rejected->code={code}')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {trade.pnl:.2f}, NET {trade.pnlcomm:.2f}')

    def stop(self):
        print('Strategy Finish!!!')

        # with open(r'C:\Users\AndrewX\Desktop\cache_log.txt', 'w') as f:
        #     res = '\n'.join(self.cache_log)
        #     f.write(res)
        doprint = (False, True)[0]

        if doprint:
            daily_pnl_df = pd.DataFrame(self.daily_pnl.items(), columns=['DATE_T', 'PNL'])
            append_df = pd.DataFrame([[101, round((self.broker.getvalue() / init_cash - 1) * 100, 2)]],
                                     columns=['DATE_T', 'PNL'])
            daily_pnl_df = daily_pnl_df.append(append_df, ignore_index=True)
            trade_df = pd.DataFrame(self.trade_info,
                                    columns=['DATE_T', 'STK_CODE', 'ACTION', 'PRICE', 'AMOUNT', 'COMM'])
            pnl_table = 'TEST_DAILY_PNL'
            trade_table = 'TEST_TRADE_INFO'
            # daily_pnl_df = daily_pnl_df.loc[(daily_pnl_df['DATE_T'] >= 20230428) | (daily_pnl_df['DATE_T'] == 101)]
            insert_db(daily_pnl_df, pnl_table)
            insert_db(trade_df, trade_table)
            print(self.trail_dict)
            print('insert table ok')


init_cash = 1e6
res = get_trade_info()
trade_code = res['trade_code']
# stocks = res['stocks']
index = '399300'
stocks = ['510050', '510300', '510500', '159915', '588000', '159928', '512010', '159939',
          '510230', '562910', '513500', '513100', '515030', '159995', '512660', '515790', '512790']
stocks.insert(0, index)
# stocks = stocks[:5]


def run_backtrade():
    cerebro = bt.Cerebro()
    cerebro.addstrategy(TKH_KDL)
    data_all = get_data(stocks, source='db')
    index = stocks[0]
    data0 = data_all[index]
    datafeed0 = bt.feeds.PandasData(dataname=data0)
    cerebro.adddata(datafeed0, name=f'{index}')

    date_data = set_date()
    st_dt = date_data['st_dt']
    et_dt = date_data['et_dt']

    for i, code in enumerate(data_all):
        if code == index:
            df = data_all[code]
            df = df.loc[:, ['close']]
            df.rename(columns={'close': '399300'}, inplace=True)
            continue
        data = data_all[code]
        datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_dt, todate=et_dt)
        cerebro.adddata(datafeed, name=code)
        print(f'{code} feeds ok')

    # 初始资金 1,000,000.0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    # cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    strat = result[0]
    # st = strat.params.st
    pyfoliozer = strat.analyzers.getbyname('pyfolio')
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    # print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

    returns, positions, transactions, gross_lev = pyfoliozer.get_pf_items()
    print(st_dt.date())
    returns = returns[returns.index.date >= st_dt.date()]
    returns.index = returns.index.date
    returns.index = pd.to_datetime(returns.index)

    print(f'returns={returns}')
    html_file = 'test.html'
    qs.reports.html(returns, benchmark=df, title='grid', output=html_file, download_filename=html_file)
    # qs.reports.html(returns, title='grid', output=html_file, download_filename=html_file)
    qs.reports.basic(returns)
    f = webbrowser.open(html_file)


def main():
    run_backtrade()


if __name__ == '__main__':
    main()
